Greptile Summary

This PR changes Blackwell backend selection so THD training with dropout prefers FlashAttention 2 over the slower FusedAttention path.

• Adds a final-selection override for SM100/SM103 when FA2 and FusedAttention are both available.
• Logs the FusedAttention preference override after backend availability has been finalized.
• Adds a Blackwell-specific test covering the resulting backend choice.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the previously reported advisory-ordering issue is resolved by placing the preference and its debug message after backend availability is finalized.

Important Files Changed

_{Reviews (3): Last reviewed commit: "Keep the THD dropout condition under the..." | Re-trigger Greptile}

Thanks for the detailed reporting in #3312!

I agree that this affects users' performance and without warning, but I don't think it's TE's responsibility to cache these default-value issues coming from Megatron. Given that dropout is going out of fashion, Megatron can probably change their default to 0 but in a gradual deprecation code cycle. This way, it wouldn't incorrectly and silently call FlashAttention for dropout either.

With cuDNN, we can check with them to see if they can speed up the dropout implementation on Blackwell for THD, to match FlashAttention's performance. For now, we can disable FusedAttention for THD + dropout + Blackwell, and allow users to use FlashAttention instead. @KshitijLakhani, could you please lead the discussion with cuDNN, file a bug if necessary, and guide @bzantium to disable FusedAttention in this PR? Thanks!

Thanks for the detailed reporting in #3312!

I agree that this affects users' performance and without warning, but I don't think it's TE's responsibility to cache these default-value issues coming from Megatron. Given that dropout is going out of fashion, Megatron can probably change their default to 0 but in a gradual deprecation code cycle. This way, it wouldn't incorrectly and silently call FlashAttention for dropout either.

With cuDNN, we can check with them to see if they can speed up the dropout implementation on Blackwell for THD, to match FlashAttention's performance. For now, we can disable FusedAttention for THD + dropout + Blackwell, and allow users to use FlashAttention instead. @KshitijLakhani, could you please lead the discussion with cuDNN, file a bug if necessary, and guide @bzantium to disable FusedAttention in this PR? Thanks!

@bzantium thanks for reporting this.
I definitely am able to reproduce your observations with TE 2.17 as well, so let me take up this discussion with the cuDNN team and report back once I have a good solution.
In the mean time, let me recommend a simple, no code changes WAR.

Please set NVTE_FUSED_ATTN=0 NVTE_FLASH_ATTN=1 (hopefully this unblocks you temporarily)
I independently validated this on GB200/SM100 with TE 2.17, cuDNN 9.24, and FA2 2.7.4:

• THD FusedAttention, dropout 0.0: 0.626 ms
• THD FusedAttention, dropout 0.1: 4.407 ms — 7.04× slower
• THD FA2, dropout 0.0: 0.695 ms
• THD FA2, dropout 0.1: 0.728 ms — only 1.05× slower

For THD with dropout 0.1, switching to FA2 reduced time from 4.407 ms to approximately 0.728 ms, about a 6.1× speedup.
This agrees with your B300 results:

• FusedAttention: 0.585 → 3.060 ms, or 5.24×
• FA2 dropout overhead: only approximately +0.024 ms

To verify that FA2 is selected, temporarily use:
NVTE_FUSED_ATTN=0 NVTE_FLASH_ATTN=1 NVTE_UNFUSED_ATTN=0 NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2

The log should report something like this:
Selected backend = FlashAttention (2.7.4.post1)

NOTE: I am suggesting NVTE_UNFUSED_ATTN=0 only as a validation guard to prevent silent fallback.

@bzantium Here's a quick update: I discussed with the cuDNN team and it's a known shortcoming. The issue is not unified v/s composite but it is the slow dropout kernels used for this very specific combination being discussed here. The fact that the slow dropout kernels are being triggered in a composite setup worsens the timing. cuDNN does not have this on their roadmap for the foreseeable future as it is low priority and the suggested "fix" is for TE users to use FA2 instead if performance is of importance.

If you'd like to contribute to TE, which we highly encourage, I've reviewed your PR and added suggested changes. Please review these changes and commit them. Once you've done this I can help you launch CI on this PR and review/approve the PR to have it merged in to main.

On the contrary, if you'd rather have me add in these changes, I'm happy to do so in a separate PR. Let me know.

Notes:

• Please make your branch up-to-date with main
• Please change the title of the PR to something more suitable now that we know it isn't a composite engine problem per se. A suggestion is Prefer Flash Attn 2 instead of Fused Attn for THD dropout on Blackwell but feel free to use your own

Lastly, here's a small mock up test I came up with using my agent. I just asked it to write a test to confirm the backend choice. This is juts a guidance but it does the trick. Feel free to be creative :)

  @pytest.mark.skipif(
      device_compute_capability not in ((10, 0), (10, 3)),
      reason="This backend preference applies only to SM100/SM103.",
  )
  @pytest.mark.skipif(
      not FlashAttentionUtils.is_installed,
      reason="A supported FlashAttention 2 installation is required.",
  )
  @pytest.mark.skipif(
      get_cudnn_version() < (9, 0, 0),
      reason="cuDNN 9.0.0+ is required for THD FusedAttention.",
  )
  def test_thd_dropout_prefers_flash_attention_2(monkeypatch):
      # Ensure both candidate backends are enabled for this test. Pytest restores
      # the previous environment automatically after the test completes.
      monkeypatch.setenv("NVTE_FLASH_ATTN", "1")
      monkeypatch.setenv("NVTE_FLASH_ATTN_V2", "1")
      monkeypatch.setenv("NVTE_FUSED_ATTN", "1")
      attention_params = AttentionParams(
          qkv_dtype=torch.bfloat16,
          qkv_layout="thd_thd_thd",
          batch_size=4,
          num_heads=16,
          num_gqa_groups=16,
          max_seqlen_q=1024,
          max_seqlen_kv=1024,
          head_dim_qk=128,
          head_dim_v=128,
          attn_mask_type="padding_causal",
          core_attention_bias_shape=None,
          attention_dropout=0.1,
          is_training=True,
      )
      (
          use_flash_attention,
          flash_attention_backend,
          use_fused_attention,
          fused_attention_backend,
          _,
          available_backends,
      ) = get_attention_backend(attention_params)
      assert available_backends[0]  # Real FA2 support check
      assert available_backends[1]  # Real cuDNN FusedAttention support check
      assert use_flash_attention
      assert flash_attention_backend == FlashAttentionUtils.version
      assert not use_fused_attention
      assert fused_attention_backend is None

@bzantium just circling back to check in if you were able to try the suggested changes ?
Please let me know if you'd prefer me to create a separate PR to fix the issue. Thanks !

@KshitijLakhani thanks for chasing this with the cuDNN team, and sorry for the slow reply.

Took both suggestions in 54f463b. The advisory is gone, and the FA2 preference now sits after use_flash_attention_2 is finalized so FusedAttention only gets disabled when FA2 is really usable. Rebased on main and retitled too.

For the test I stayed close to your mock-up but added assert available_backends[0] and [1] up front, otherwise a box without FA2 or cuDNN would pass for the wrong reason. Kept it in test_attention.py since the QA scripts name test files one by one.

Ran it on an SM103 box with both backends available: FusedAttention (NVTE_F16_arbitrary_seqlen) before, FA2 after. Happy for you to kick off CI whenever.